import cv2
import os
import sys

import numpy as np
import torch
from basicsr.utils import img2tensor, tensor2img
from basicsr.utils.download_util import load_file_from_url
from torchvision.transforms.functional import normalize
from facexlib.utils.face_restoration_helper import FaceRestoreHelper

if __name__ == '__main__':
    sys.path.append('../')
from gfpgan.archs.gfpgan_bilinear_arch import GFPGANBilinear
from gfpgan.archs.gfpganv1_arch import GFPGANv1
from gfpgan.archs.gfpganv1_clean_arch import GFPGANv1Clean

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))


class GFPGANFaceRestorer():
    """Helper for restoration with GFPGAN.

    It will detect and crop faces, and then resize the faces to 512x512.
    GFPGAN is used to restored the resized faces.
    The background is upsampled with the bg_upsampler.
    Finally, the faces will be pasted back to the upsample background image.

    Args:
        model_path (str): The path to the GFPGAN model. It can be urls (will first download it automatically).
        upscale (float): The upscale of the final output. Default: 2.
        arch (str): The GFPGAN architecture. Option: clean | original. Default: clean.
        channel_multiplier (int): Channel multiplier for large networks of StyleGAN2. Default: 2.
        bg_upsampler (nn.Module): The upsampler for the background. Default: None.
    """

    def __init__(self, model_path=os.path.join(ROOT_DIR, '../experiments/pretrained_models', 'GFPGANv1.4.pth'),
                 upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=None, device=None):
        self.upscale = upscale
        self.bg_upsampler = bg_upsampler

        # initialize model
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
        # initialize the GFP-GAN
        if arch == 'clean':
            self.gfpgan = GFPGANv1Clean(
                out_size=512,
                num_style_feat=512,
                channel_multiplier=channel_multiplier,
                decoder_load_path=None,
                fix_decoder=False,
                num_mlp=8,
                input_is_latent=True,
                different_w=True,
                narrow=1,
                sft_half=True)
        elif arch == 'bilinear':
            self.gfpgan = GFPGANBilinear(
                out_size=512,
                num_style_feat=512,
                channel_multiplier=channel_multiplier,
                decoder_load_path=None,
                fix_decoder=False,
                num_mlp=8,
                input_is_latent=True,
                different_w=True,
                narrow=1,
                sft_half=True)
        elif arch == 'original':
            self.gfpgan = GFPGANv1(
                out_size=512,
                num_style_feat=512,
                channel_multiplier=channel_multiplier,
                decoder_load_path=None,
                fix_decoder=True,
                num_mlp=8,
                input_is_latent=True,
                different_w=True,
                narrow=1,
                sft_half=True)
        elif arch == 'RestoreFormer':
            from gfpgan.archs.restoreformer_arch import RestoreFormer
            self.gfpgan = RestoreFormer()
        # initialize face helper
        self.face_helper = FaceRestoreHelper(
            upscale,
            face_size=512,
            crop_ratio=(1, 1),
            det_model='retinaface_resnet50',
            save_ext='png',
            use_parse=True,
            device=self.device,
            model_rootpath=os.path.join(ROOT_DIR, 'weights'))

        if model_path.startswith('https://'):
            model_path = load_file_from_url(
                url=model_path, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
        loadnet = torch.load(model_path)
        if 'params_ema' in loadnet:
            keyname = 'params_ema'
        else:
            keyname = 'params'
        self.gfpgan.load_state_dict(loadnet[keyname], strict=True)
        self.gfpgan.eval()
        self.gfpgan = self.gfpgan.to(self.device)

    @torch.no_grad()
    def enhance_face(self, img, only_center_face=False, weight=0.5):
        img_orgin_size = img.shape[:-1]  # h,w
        img_orgin_size = (img_orgin_size[-1], img_orgin_size[0])  # h,w -> w,h
        result = []
        self.face_helper.clean_all()
        self.face_helper.read_image(img)
        # get face landmarks for each face
        self.face_helper.get_face_landmarks_5(only_center_face=only_center_face, eye_dist_threshold=5)
        # eye_dist_threshold=5: skip faces whose eye distance is smaller than 5 pixels
        # TODO: even with eye_dist_threshold, it will still introduce wrong detections and restorations.
        # align and warp each face
        self.face_helper.align_warp_face()
        # face restoration
        for cropped_face, affine_matrix in zip(self.face_helper.cropped_faces, self.face_helper.affine_matrices):
            # prepare data
            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
            cropped_face_t = cropped_face_t.unsqueeze(0).to(self.device)
            try:
                output = self.gfpgan(cropped_face_t, return_rgb=False, weight=weight)[0]
                # convert to image
                restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
            except RuntimeError as error:
                print(f'\tFailed inference for GFPGAN: {error}.')
                restored_face = cropped_face
            restored_face = restored_face.astype('uint8')
            # 还原图像
            inverse_affine_matrix = cv2.invertAffineTransform(affine_matrix)
            # cv2.imwrite('restored_face_rawsize.jpg', restored_face)
            restored_face = cv2.warpAffine(restored_face, inverse_affine_matrix, img_orgin_size)
            # 考虑使用parsenet
            result.append(restored_face)
        return result

    def net_parse(self, restored_face, inv_restored, inverse_affine, w, h):
        # inference
        face_input = cv2.resize(restored_face, (512, 512), interpolation=cv2.INTER_LINEAR)
        face_input = img2tensor(face_input.astype('float32') / 255., bgr2rgb=True, float32=True)
        normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
        face_input = torch.unsqueeze(face_input, 0).to(self.device)
        with torch.no_grad():
            out = self.face_helper.face_parse(face_input)[0]
        out = out.argmax(dim=1).squeeze().cpu().numpy()

        mask = np.zeros(out.shape)
        MASK_COLORMAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 0, 0, 0]
        for idx, color in enumerate(MASK_COLORMAP):
            mask[out == idx] = color
        #  blur the mask
        mask = cv2.GaussianBlur(mask, (101, 101), 11)
        mask = cv2.GaussianBlur(mask, (101, 101), 11)
        # remove the black borders
        thres = 10
        mask[:thres, :] = 0
        mask[-thres:, :] = 0
        mask[:, :thres] = 0
        mask[:, -thres:] = 0
        mask = mask / 255.

        mask = cv2.resize(mask, restored_face.shape[:2])
        mask = cv2.warpAffine(mask, inverse_affine, (w, h), flags=3)
        inv_soft_mask = mask[:, :, None]
        pasted_face = inv_restored
        return inv_soft_mask, pasted_face

    @torch.no_grad()
    def enhance_face_2(self, origin_img, only_center_face=False, weight=0.5):
        h, w, _ = origin_img.shape
        result = []
        self.face_helper.clean_all()
        self.face_helper.read_image(origin_img)
        # get face landmarks for each face
        self.face_helper.get_face_landmarks_5(only_center_face=only_center_face, eye_dist_threshold=5)
        # eye_dist_threshold=5: skip faces whose eye distance is smaller than 5 pixels
        # TODO: even with eye_dist_threshold, it will still introduce wrong detections and restorations.
        # align and warp each face
        self.face_helper.align_warp_face()
        # face restoration
        for cropped_face, affine_matrix in zip(self.face_helper.cropped_faces, self.face_helper.affine_matrices):
            # prepare data
            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
            cropped_face_t = cropped_face_t.unsqueeze(0).to(self.device)
            try:
                output = self.gfpgan(cropped_face_t, return_rgb=False, weight=weight)[0]
                # convert to image
                restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
            except RuntimeError as error:
                print(f'\tFailed inference for GFPGAN: {error}.')
                restored_face = cropped_face
            restored_face = restored_face.astype('uint8')
            # 还原图像
            inverse_affine = cv2.invertAffineTransform(affine_matrix)
            inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w, h))
            inv_soft_mask, pasted_face = self.net_parse(restored_face, inv_restored, inverse_affine, w, h)
            final_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * origin_img
            if np.max(final_img) > 256:  # 16-bit image
                final_img = final_img.astype(np.uint16)
            else:
                final_img = final_img.astype(np.uint8)
            result.append(final_img)
        return result

    def square_parse(self, origin_img, inv_restored, inverse_affine, w, h):
        mask = np.ones(self.face_helper.face_size, dtype=np.float32)
        inv_mask = cv2.warpAffine(mask, inverse_affine, (w, h))
        # remove the black borders
        inv_mask_erosion = cv2.erode(
            inv_mask, np.ones((int(2 * self.face_helper.upscale_factor), int(2 * self.face_helper.upscale_factor)), np.uint8))
        pasted_face = inv_mask_erosion[:, :, None] * inv_restored
        total_face_area = np.sum(inv_mask_erosion)  # // 3
        # compute the fusion edge based on the area of face
        w_edge = int(total_face_area ** 0.5) // 20
        erosion_radius = w_edge * 2
        inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
        blur_size = w_edge * 2
        inv_soft_mask = cv2.GaussianBlur(inv_mask_center, (blur_size + 1, blur_size + 1), 0)
        if len(origin_img.shape) == 2:  # upsample_img is gray image
            origin_img = origin_img[:, :, None]
        inv_soft_mask = inv_soft_mask[:, :, None]
        return inv_soft_mask, pasted_face, origin_img

    @torch.no_grad()
    def enhance_face_3(self, origin_img, only_center_face=False, weight=0.5):
        h, w, _ = origin_img.shape
        result = []
        self.face_helper.clean_all()
        self.face_helper.read_image(origin_img)
        # get face landmarks for each face
        self.face_helper.get_face_landmarks_5(only_center_face=only_center_face, eye_dist_threshold=5)
        # eye_dist_threshold=5: skip faces whose eye distance is smaller than 5 pixels
        # TODO: even with eye_dist_threshold, it will still introduce wrong detections and restorations.
        # align and warp each face
        self.face_helper.align_warp_face()
        # face restoration
        for cropped_face, affine_matrix in zip(self.face_helper.cropped_faces, self.face_helper.affine_matrices):
            # prepare data
            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
            cropped_face_t = cropped_face_t.unsqueeze(0).to(self.device)
            try:
                output = self.gfpgan(cropped_face_t, return_rgb=False, weight=weight)[0]
                # convert to image
                restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(-1, 1))
            except RuntimeError as error:
                print(f'\tFailed inference for GFPGAN: {error}.')
                restored_face = cropped_face
            restored_face = restored_face.astype('uint8')
            # 还原图像
            inverse_affine = cv2.invertAffineTransform(affine_matrix)
            inv_restored = cv2.warpAffine(restored_face, inverse_affine, (w, h))
            inv_soft_mask, pasted_face, origin_img = self.square_parse(origin_img, inv_restored, inverse_affine, w, h)
            final_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * origin_img
            if np.max(final_img) > 256:  # 16-bit image
                final_img = final_img.astype(np.uint16)
            else:
                final_img = final_img.astype(np.uint8)
            result.append(final_img)
        return result


if __name__ == '__main__':
    restorer = GFPGANFaceRestorer(upscale=0)
    input_img = cv2.imread('face_resize.jpg', cv2.IMREAD_COLOR)  # face_resize  cropped_face
    restored_faces = restorer.enhance_face(
        input_img,
        only_center_face=True,
    )
    cv2.imwrite('restored_face.jpg', restored_faces[0])
    print('over!')

    pass
